#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2023-2025. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================
import os
import numpy as np

# ---------------------------------数据类型映射-----------------------------------

dtype_emu = {
    np.float16: 0,
    np.float32: 1,
    np.int16: 2,
    np.int32: 3
}

# 输出 float 固定参数
BLOCK_SIZE = 32          
OUTPUT_DTYPE_SIZE = 4   
ALIGN_NUM = BLOCK_SIZE // OUTPUT_DTYPE_SIZE 
UB_BLOCK_NUM = 100       
BUFFER_NUM = 2           
MAX_AVAILABLE_UB_BLOCK_NUM = (UB_BLOCK_NUM // BUFFER_NUM) * BUFFER_NUM  

# -------------------------------------------------------------------------------
def tiling_params_calc(length: int, align_num: int) -> tuple:
    
    max_per_core_elem = align_num * MAX_AVAILABLE_UB_BLOCK_NUM  
    tile_num = length // max_per_core_elem

    if (length % max_per_core_elem == 0) or (tile_num == 0):
        if tile_num == 0:
            tile_num = 1 
        if length < max_per_core_elem:

            block_count = (length + align_num - 1) // align_num 
            block_count = (block_count + BUFFER_NUM - 1) // BUFFER_NUM * BUFFER_NUM
            tile_length = block_count * align_num
            last_tile_length = tile_length
        else:
            tile_length = max_per_core_elem
            last_tile_length = length - (tile_num - 1) * tile_length
    else:
        tile_num += 1
        tile_length = max_per_core_elem
        last_tile_length = length - (tile_num - 1) * tile_length

    return tile_num, tile_length, last_tile_length


# --------------------------------------------------------------------------------
def generate_tiling_data(num: int, input_dtype: np.dtype, block_dim: int) -> np.ndarray:

    tiling_fields = np.zeros(18, dtype=np.uint32)  
    total_length_aligned = num if (num % ALIGN_NUM == 0) else ((num + ALIGN_NUM - 1) // ALIGN_NUM) * ALIGN_NUM
    
    tiling_fields[0] = num
    tiling_fields[1] = dtype_emu[input_dtype]
    tiling_fields[2] = total_length_aligned
    
    total_block_count = total_length_aligned // ALIGN_NUM 

    if (total_block_count % block_dim) == 0:
        tiling_fields[17] = 1

        block_length = total_length_aligned // block_dim
        tile_num, tile_length, last_tile_length = tiling_params_calc(block_length, ALIGN_NUM)
        
        tiling_fields[3] = block_length
        tiling_fields[4] = tile_num
        tiling_fields[5] = tile_length
        tiling_fields[6] = last_tile_length
    
    else:
        tiling_fields[17] = 0 
        
        former_num = total_block_count % block_dim 
        tail_num = block_dim - former_num    
        former_block_count = (total_block_count + block_dim - 1) // block_dim
        former_length = former_block_count * ALIGN_NUM
        tail_block_count = total_block_count // block_dim
        tail_length = tail_block_count * ALIGN_NUM
        former_tile_num, former_tile_len, former_last_tile_len = tiling_params_calc(former_length, ALIGN_NUM)
        tail_tile_num, tail_tile_len, tail_last_tile_len = tiling_params_calc(tail_length, ALIGN_NUM)
        
        tiling_fields[7] = former_num
        tiling_fields[8] = former_length            
        tiling_fields[9] = former_tile_num         
        tiling_fields[10] = former_tile_len     
        tiling_fields[11] = former_last_tile_len
        tiling_fields[12] = tail_num               
        tiling_fields[13] = tail_length              
        tiling_fields[14] = tail_tile_num        
        tiling_fields[15] = tail_tile_len          
        tiling_fields[16] = tail_last_tile_len     
    
    return tiling_fields
# ----------------------------------------------------------------------------------
def linspace_custom(start, end, num, output_dtype=np.float32):
    """自定义 linspace：输入 start/end 支持多类型，输出 float32"""
    if num <= 1:
        return np.array([start], dtype=output_dtype)
    start_f32 = start.astype(np.float32)
    end_f32 = end.astype(np.float32)
    step = (end_f32 - start_f32) / (num - 1)
    seq = start_f32 + np.arange(num, dtype=np.float32) * step
    return seq.astype(output_dtype)
# -------------------------- ---------------------------------------------------------
def gen_golden_data_simple():
    np.random.seed()  
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)
    
    # ==============================配置参数================================
    # input_dtype = np.float16    # 对应 SIG_HALF=0
    input_dtype = np.float32   # 对应 SIG_FLOAT=1
    # input_dtype = np.int16     # 对应 SIG_INT16=2
    # input_dtype = np.int32     # 对应 SIG_INT32=3
    output_dtype = np.float32  
    ns_dtype = np.uint32       
    
    # 测试场景：
    # - 核间均分：block_dim=8, num=1024（1024//8=128，128//8=16 均分）
    # - 核间非均分：block_dim=17, num=1023（1023对齐后1024，1024//8=128，128%17≠0）
    input_shape = [8, 1023]  # [核心数, 冗余维度]
    block_dim = input_shape[0]
    
    # 参数取值范围
    min_n = 1                       # num 最小值
    max_n = 1024                    # num 最大值
    start_range = (-10.0, 10.0)     # start 取值范围
    end_offset_range = (0.1, 20.0)  # end - start 的差值范围

    # ==============================生成数据 ================================
    starts = np.random.uniform(*start_range, size=1).astype(input_dtype)
    end_offset = np.random.uniform(*end_offset_range, size=1).astype(input_dtype)
    ends = starts + end_offset
    ns = np.random.randint(min_n, max_n + 1, size=1, dtype=ns_dtype)
    num = ns[0]  
    golden = linspace_custom(starts[0], ends[0], num, output_dtype=output_dtype)
    tiling_data = generate_tiling_data(num, input_dtype, block_dim)
    # ========================== 写入二进制文件 =============================
    starts.tofile("./input/starts.bin")
    ends.tofile("./input/ends.bin")
    ns.tofile("./input/ns.bin")
    tiling_data.tofile("./input/input_tiling.bin")
    golden.tofile("./output/golden.bin")

if __name__ == "__main__":
    gen_golden_data_simple()